/*
 * Copyright 2013-2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.cloud.gateway.filter;

import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import reactor.core.publisher.Mono;

import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.gateway.filter.headers.HttpHeadersFilter;
import org.springframework.core.Ordered;
import org.springframework.http.HttpHeaders;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.socket.CloseStatus;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import org.springframework.web.reactive.socket.client.WebSocketClient;
import org.springframework.web.reactive.socket.server.WebSocketService;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.util.UriComponentsBuilder;

import static org.springframework.cloud.gateway.filter.headers.HttpHeadersFilter.filterRequest;
import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;
import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.PRESERVE_HOST_HEADER_ATTRIBUTE;
import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.containsEncodedParts;
import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.isAlreadyRouted;
import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setAlreadyRouted;

/**
 * @author Spencer Gibb
 * @author Nikita Konev
 * <p>
 *     WebsocketRoutingFilter 是 Spring Cloud Gateway 中用于处理 WebSocket 请求的核心过滤器。
 *     它的主要功能是将 HTTP 或 HTTPS 请求转换为 WebSocket（ws 或 wss）请求，
 *     并通过代理的方式将客户端的 WebSocket 通信转发到后端服务。以下是对该类的详细分析和说明。
 *
 *     WebsocketRoutingFilter 的主要职责是识别 WebSocket 握手请求，并将其路由到后端服务。
 * 		它通过检查请求的 Upgrade 头部是否为 WebSocket 来判断是否需要处理 WebSocket 请求。
 *
 * 	1. 示例场景
 * 		假设我们有一个后端 WebSocket 服务运行在 ws://localhost:8081/websocket，
 * 		我们需要通过 Spring Cloud Gateway 将客户端的 WebSocket 请求路由到该服务。
 * 	2. 项目结构
 * 		2.1 添加依赖
 * 		在 pom.xml 中添加以下依赖：
		 * <dependencies>
		 *     <!-- Spring Boot Starter for WebFlux -->
		 *     <dependency>
		 *         <groupId>org.springframework.boot</groupId>
		 *         <artifactId>spring-boot-starter-webflux</artifactId>
		 *     </dependency>
		 *
		 *     <!-- Spring Cloud Gateway -->
		 *     <dependency>
		 *         <groupId>org.springframework.cloud</groupId>
		 *         <artifactId>spring-cloud-starter-gateway</artifactId>
		 *     </dependency>
		 *
		 *     <!-- Reactor Netty for WebSocket support -->
		 *     <dependency>
		 *         <groupId>io.projectreactor.netty</groupId>
		 *         <artifactId>reactor-netty</artifactId>
		 *     </dependency>
		 * </dependencies>
 *
 *	 2.2 配置网关路由
 *	 在 application.yml 中配置路由规则，将 WebSocket 请求转发到后端服务。
	 * spring:
	 *   cloud:
	 *     gateway:
	 *       routes:
	 *         - id: websocket_route
	 *           uri: ws://localhost:8081 # 后端 WebSocket 服务地址
	 *           predicates:
	 *             - Path=/websocket # 匹配路径
	 *           filters:
	 *             - name: WebSocketRoutingFilter # 使用 WebSocket 路由过滤器
 *  2.3 创建主类
	 * 创建一个 Spring Boot 应用程序入口类。
	 * import org.springframework.boot.SpringApplication;
	 * import org.springframework.boot.autoconfigure.SpringBootApplication;
	 * @SpringBootApplication
	 * public class GatewayApplication {
	 *     public static void main(String[] args) {
	 *         SpringApplication.run(GatewayApplication.class, args);
	 *     }
	 * }
 *
 * 2.4 测试后端 WebSocket 服务
 * 为了测试，我们可以创建一个简单的 WebSocket 服务。以下是后端服务的代码示例：
 *
 * import org.springframework.context.annotation.Configuration;
 * import org.springframework.web.socket.config.annotation.EnableWebSocket;
 * import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
 * import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
 * import org.springframework.web.socket.handler.TextWebSocketHandler;
 *
 * @Configuration
 * @EnableWebSocket
 * public class WebSocketConfig implements WebSocketConfigurer {
 *
 *     @Override
 *     public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
 *         registry.addHandler(new EchoWebSocketHandler(), "/websocket").setAllowedOrigins("*");
 *     }
 *
 *     static class EchoWebSocketHandler extends TextWebSocketHandler {
 *         @Override
 *         protected void handleTextMessage(WebSocketSession session, org.springframework.web.socket.TextMessage message) throws Exception {
 *             String payload = message.getPayload();
 *             session.sendMessage(new org.springframework.web.socket.TextMessage("Echo: " + payload));
 *         }
 *     }
 * }
 * 启动后端服务，确保它运行在 ws://localhost:8081/websocket。
 * </p>
 */
public class WebsocketRoutingFilter implements GlobalFilter, Ordered {

	/**
	 * Sec-Websocket protocol.
	 */
	public static final String SEC_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol";

	private static final Log log = LogFactory.getLog(WebsocketRoutingFilter.class);

	private final WebSocketClient webSocketClient;

	private final WebSocketService webSocketService;

	private final ObjectProvider<List<HttpHeadersFilter>> headersFiltersProvider;

	// do not use this headersFilters directly, use getHeadersFilters() instead.
	private volatile List<HttpHeadersFilter> headersFilters;

	public WebsocketRoutingFilter(WebSocketClient webSocketClient, WebSocketService webSocketService,
			ObjectProvider<List<HttpHeadersFilter>> headersFiltersProvider) {
		this.webSocketClient = webSocketClient;
		this.webSocketService = webSocketService;
		this.headersFiltersProvider = headersFiltersProvider;
	}

	/**
	 * 将 HTTP 或 HTTPS 请求转换为 WebSocket 请求（ws 或 wss）。
	 * 通过 convertHttpToWs 方法实现协议转换
	 * @param scheme
	 * @return
	 */
	/* for testing */
	static String convertHttpToWs(String scheme) {
		scheme = scheme.toLowerCase();
		return "http".equals(scheme) ? "ws" : "https".equals(scheme) ? "wss" : scheme;
	}

	@Override
	public int getOrder() {
		// Before NettyRoutingFilter since this routes certain http requests
		return Ordered.LOWEST_PRECEDENCE - 1;
	}

	/**
	 * 此函数是Spring Cloud Gateway项目中的一个自定义过滤器实现，用于处理WebSocket协议的路由和消息代理。在Spring Cloud Gateway中，过滤器是处理请求和响应的关键组件，可以在请求被发送到下游服务之前或之后对请求和响应进行修改。该函数通过实现GatewayFilter接口的filter方法，参与到网关的请求处理流程中。
	 * 首先，函数通过调用changeSchemeIfIsWebSocketUpgrade方法，检查当前请求是否是一个WebSocket升级请求，并在必要时更改请求的scheme。这是因为在HTTP协议中，WebSocket连接通常是通过HTTP升级请求来建立的。
	 * 接着，函数从ServerWebExchange对象中获取请求URL，并检查其scheme是否为ws（WebSocket非加密）或wss（WebSocket加密）。这是为了确定请求是否需要使用WebSocket协议进行处理。
	 * 如果请求已经被路由过（通过isAlreadyRouted方法检查），或者请求的scheme不是ws或wss，则函数直接调用chain.filter(exchange)，将请求传递给下一个过滤器或最终的路由处理器，不进行WebSocket处理。
	 * 如果请求需要使用WebSocket协议，函数会先通过setAlreadyRouted方法标记请求为已路由，以避免重复处理。然后，获取请求头信息，并通过filterRequest方法应用一系列请求头过滤器，对请求头进行修改或清理。
	 * 接下来，函数通过getProtocols方法从请求头中提取WebSocket子协议列表。这些子协议是在WebSocket握手过程中由客户端提供的，用于指定客户端和服务端之间通信使用的特定协议。
	 * 最后，函数使用webSocketService.handleRequest方法处理WebSocket请求。该方法接受一个ServerWebExchange对象和一个ProxyWebSocketHandler实例作为参数。ProxyWebSocketHandler是一个实现了WebSocket消息代理的类，负责在客户端和后端服务之间转发WebSocket消息。它使用提供的请求URL、WebSocket客户端、过滤后的请求头和子协议列表来建立WebSocket连接，并处理消息的转发。
	 * @param exchange the current server exchange
	 * @param chain provides a way to delegate to the next filter
	 * @return
	 */
	@Override
	public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
		//1、首先，函数通过调用changeSchemeIfIsWebSocketUpgrade方法，检查当前请求是否是一个WebSocket升级请求，并在必要时更改请求的scheme。这是因为在HTTP协议中，WebSocket连接通常是通过HTTP升级请求来建立的。
		changeSchemeIfIsWebSocketUpgrade(exchange);
		//2、接着，函数从ServerWebExchange对象中获取请求URL，并检查其scheme是否为ws（WebSocket非加密）或wss（WebSocket加密）。这是为了确定请求是否需要使用WebSocket协议进行处理。
		URI requestUrl = exchange.getRequiredAttribute(GATEWAY_REQUEST_URL_ATTR);
		String scheme = requestUrl.getScheme();

		if (isAlreadyRouted(exchange) || (!"ws".equals(scheme) && !"wss".equals(scheme))) {
			return chain.filter(exchange);
		}
		setAlreadyRouted(exchange);

		HttpHeaders headers = exchange.getRequest().getHeaders();
		HttpHeaders filtered = filterRequest(getHeadersFilters(), exchange);

		List<String> protocols = getProtocols(headers);

		//使用 ProxyWebSocketHandler 实现客户端与后端服务之间的 WebSocket 通信代理。
		//它负责在客户端和后端服务之间转发 WebSocket 消息。
		return this.webSocketService.handleRequest(exchange,
				new ProxyWebSocketHandler(requestUrl, this.webSocketClient, filtered, protocols));
	}

	/* for testing */ List<String> getProtocols(HttpHeaders headers) {
		List<String> protocols = headers.get(SEC_WEBSOCKET_PROTOCOL);
		if (protocols != null) {
			ArrayList<String> updatedProtocols = new ArrayList<>();
			for (int i = 0; i < protocols.size(); i++) {
				String protocol = protocols.get(i);
				updatedProtocols.addAll(Arrays.asList(StringUtils.tokenizeToStringArray(protocol, ",")));
			}
			protocols = updatedProtocols;
		}
		return protocols;
	}

	/* for testing */ List<HttpHeadersFilter> getHeadersFilters() {
		if (this.headersFilters == null) {
			this.headersFilters = this.headersFiltersProvider.getIfAvailable(ArrayList::new);

			// remove host header unless specifically asked not to
			headersFilters.add((headers, exchange) -> {
				HttpHeaders filtered = new HttpHeaders();
				filtered.addAll(headers);
				filtered.remove(HttpHeaders.HOST);
				boolean preserveHost = exchange.getAttributeOrDefault(PRESERVE_HOST_HEADER_ATTRIBUTE, false);
				if (preserveHost) {
					String host = exchange.getRequest().getHeaders().getFirst(HttpHeaders.HOST);
					filtered.add(HttpHeaders.HOST, host);
				}
				return filtered;
			});

			headersFilters.add((headers, exchange) -> {
				HttpHeaders filtered = new HttpHeaders();
				for (Map.Entry<String, List<String>> entry : headers.entrySet()) {
					if (!entry.getKey().toLowerCase().startsWith("sec-websocket")) {
						filtered.addAll(entry.getKey(), entry.getValue());
					}
				}
				return filtered;
			});
		}

		return this.headersFilters;
	}

	/**
	 * 该函数changeSchemeIfIsWebSocketUpgrade在基于Java的Spring Cloud Gateway项目中起着关键作用，
	 * 特别是在处理WebSocket升级请求时调整URI的scheme。
	 * 以下是该函数的具体作用、代码解释以及它如何与整个项目相结合：
	 * 函数首先检查传入的ServerWebExchange对象中是否包含一个升级为WebSocket的请求。如果是，函数会进一步检查当前请求的scheme（http或https），并根据情况将其转换为WebSocket对应的scheme（ws或wss）。
	 * 这一转换是基于客户端请求中携带的Upgrade头部字段，表明客户端希望升级到WebSocket连接。完成scheme转换后，函数会构建一个新的URI，并将其放回到ServerWebExchange对象的属性中，以便后续的过滤器或处理器可以使用这个已经调整过的URI。
	 *
	 * 详细解释代码段：
	 * URI requestUrl = exchange.getRequiredAttribute(GATEWAY_REQUEST_URL_ATTR);：
	 * 	从ServerWebExchange对象中获取表示请求URL的属性。这是一个必须的属性，因此使用getRequiredAttribute方法来获取。
	 * String scheme = requestUrl.getScheme().toLowerCase();：
	 * 	获取URI的scheme部分（如http或https），并转换为小写，确保后续比较时不区分大小写。
	 * String upgrade = exchange.getRequest().getHeaders().getUpgrade();：
	 * 	从请求头中获取Upgrade字段的值，这个值指示了客户端希望升级到的协议（例如WebSocket）。
	 *
	 * if ("WebSocket".equalsIgnoreCase(upgrade) && ("http".equals(scheme) || "https".equals(scheme))) {：
	 * 	判断如果Upgrade头部是"WebSocket"，并且当前URI的scheme是http或https，则进入条件块进行scheme的转换。
	 *
	 * String wsScheme = convertHttpToWs(scheme);：
	 * 	调用convertHttpToWs方法将http或https转换为对应的WebSocket scheme（ws或wss）。
	 *
	 * boolean encoded = containsEncodedParts(requestUrl);：
	 * 	检查URL中是否包含编码部分，这个检查用于后续构建新的URI时是否需要对特定部分进行编码。
	 *
	 * URI wsRequestUrl = UriComponentsBuilder.fromUri(requestUrl).scheme(wsScheme).build(encoded).toUri();：
	 * 		使用调整后的scheme和原始的URL信息构建一个新的URI对象。
	 *
	 * exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, wsRequestUrl);：
	 * 	将新的URI对象放回到ServerWebExchange的属性中，供后续使用。
	 *
	 * if (log.isTraceEnabled()) { log.trace("changeSchemeTo:[" + wsRequestUrl + "]"); }：如果启用了trace日志级别，则记录转换后的URI。
	 *
	 * 此函数在WebSocket路由过滤器WebsocketRoutingFilter中被调用，用于处理WebSocket升级请求。在filter方法中，该函数帮助判断并转换请求URI的scheme，以确保WebSocket连接可以正确建立。例如，在测试方法testEncodedUrl中模拟了一个包含编码URL的场景，验证了函数能正确处理这种情况。
	 *
	 * 在项目中的作用，这个函数是网关处理WebSocket请求的关键步骤之一，确保了客户端可以通过HTTP/HTTPS协议升级到WebSocket连接，并且保持了请求的URL结构和参数在升级后仍然有效和正确。
	 *
	 * 整体而言，该函数是Spring Cloud Gateway处理WebSocket升级请求的重要组成部分，它通过动态调整请求的URI scheme，使得WebSocket连接可以顺利建立和处理，从而支持了更复杂的网关路由和应用场景。
	 * @param exchange
	 */
	static void changeSchemeIfIsWebSocketUpgrade(ServerWebExchange exchange) {
		// Check the Upgrade
		//从ServerWebExchange对象中获取表示请求URL的属性。这是一个必须的属性，因此使用getRequiredAttribute方法来获取。
		URI requestUrl = exchange.getRequiredAttribute(GATEWAY_REQUEST_URL_ATTR);
		//获取URI的scheme部分（如http或https），并转换为小写，确保后续比较时不区分大小写。
		String scheme = requestUrl.getScheme().toLowerCase();
		//	从请求头中获取Upgrade字段的值，这个值指示了客户端希望升级到的协议（例如WebSocket）。
		String upgrade = exchange.getRequest().getHeaders().getUpgrade();
		// change the scheme if the socket client send a "http" or "https"
		//判断如果Upgrade头部是"WebSocket"，并且当前URI的scheme是http或https，则进入条件块进行scheme的转换。
		if ("WebSocket".equalsIgnoreCase(upgrade) && ("http".equals(scheme) || "https".equals(scheme))) {
			//调用convertHttpToWs方法将http或https转换为对应的WebSocket scheme（ws或wss）。
			String wsScheme = convertHttpToWs(scheme);
			//检查URL中是否包含编码部分，这个检查用于后续构建新的URI时是否需要对特定部分进行编码。
			boolean encoded = containsEncodedParts(requestUrl);
			//使用调整后的scheme和原始的URL信息构建一个新的URI对象。
			URI wsRequestUrl = UriComponentsBuilder.fromUri(requestUrl).scheme(wsScheme).build(encoded).toUri();
			//	将新的URI对象放回到ServerWebExchange的属性中，供后续使用。
			exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, wsRequestUrl);
			if (log.isTraceEnabled()) {
				log.trace("changeSchemeTo:[" + wsRequestUrl + "]");
			}
		}
	}

	private static class ProxyWebSocketHandler implements WebSocketHandler {

		private final WebSocketClient client;

		private final URI url;

		private final HttpHeaders headers;

		private final List<String> subProtocols;

		ProxyWebSocketHandler(URI url, WebSocketClient client, HttpHeaders headers, List<String> protocols) {
			this.client = client;
			this.url = url;
			this.headers = headers;
			if (protocols != null) {
				this.subProtocols = protocols;
			}
			else {
				this.subProtocols = Collections.emptyList();
			}
		}

		@Override
		public List<String> getSubProtocols() {
			return this.subProtocols;
		}

		@Override
		public Mono<Void> handle(WebSocketSession session) {
			// pass headers along so custom headers can be sent through
			return client.execute(url, this.headers, new WebSocketHandler() {

				private CloseStatus adaptCloseStatus(CloseStatus closeStatus) {
					int code = closeStatus.getCode();
					if (code > 2999 && code < 5000) {
						return closeStatus;
					}
					switch (code) {
						case 1000:
						case 1001:
						case 1002:
						case 1003:
						case 1007:
						case 1008:
						case 1009:
						case 1010:
						case 1011:
							return closeStatus;
						case 1004:
							// Should not be used in a close frame
							// RESERVED;
						case 1005:
							// Should not be used in a close frame
							// return CloseStatus.NO_STATUS_CODE;
						case 1006:
							// Should not be used in a close frame
							// return CloseStatus.NO_CLOSE_FRAME;
						case 1012:
							// Not in RFC6455
							// return CloseStatus.SERVICE_RESTARTED;
						case 1013:
							// Not in RFC6455
							// return CloseStatus.SERVICE_OVERLOAD;
						case 1015:
							// Should not be used in a close frame
							// return CloseStatus.TLS_HANDSHAKE_FAILURE;
						default:
							return CloseStatus.PROTOCOL_ERROR;
					}
				}

				@Override
				public Mono<Void> handle(WebSocketSession proxySession) {
					Mono<Void> serverClose = proxySession.closeStatus()
						.filter(__ -> session.isOpen())
						.map(this::adaptCloseStatus)
						.flatMap(session::close);
					Mono<Void> proxyClose = session.closeStatus()
						.filter(__ -> proxySession.isOpen())
						.map(this::adaptCloseStatus)
						.flatMap(proxySession::close);
					// Use retain() for Reactor Netty
					Mono<Void> proxySessionSend = proxySession
						.send(session.receive().doOnNext(WebSocketMessage::retain).doOnNext(webSocketMessage -> {
							if (log.isTraceEnabled()) {
								log.trace("proxySession(send from client): " + proxySession.getId()
										+ ", corresponding session:" + session.getId() + ", packet: "
										+ webSocketMessage.getPayloadAsText());
							}
						}));
					// .log("proxySessionSend", Level.FINE);
					Mono<Void> serverSessionSend = session
						.send(proxySession.receive().doOnNext(WebSocketMessage::retain).doOnNext(webSocketMessage -> {
							if (log.isTraceEnabled()) {
								log.trace("session(send from backend): " + session.getId()
										+ ", corresponding proxySession:" + proxySession.getId() + " packet: "
										+ webSocketMessage.getPayloadAsText());
							}
						}));
					// .log("sessionSend", Level.FINE);
					// Ensure closeStatus from one propagates to the other
					Mono.when(serverClose, proxyClose).subscribe();
					// Complete when both sessions are done
					return Mono.zip(proxySessionSend, serverSessionSend).then();
				}

				/**
				 * Copy subProtocols so they are available downstream.
				 * @return available subProtocols.
				 */
				@Override
				public List<String> getSubProtocols() {
					return ProxyWebSocketHandler.this.subProtocols;
				}
			});
		}

	}

}
